package algorithm.ordered_table;

public class CountOfSubArrayRangeSum {

    //求在一个给定的数组中满足某个约束的子数组的数量，该约束是子数组的和在给定的范围内。
    //当出现求满足某个约束的“子数组”的数量的题目时，一般的求解思路：以某个位置作为子数组的“末尾位置”求数量

    //时间复杂度O（N*logN）
    public static int countRangeSum2(int[] arr, int lower, int up){
        long sum = 0;
        int ans = 0;
        SizeBalancedTreeMap map = new SizeBalancedTreeMap();
        map.put(0);
        for (int i = 0; i < arr.length; i++) {
            sum += arr[i];
            long num1 = map.lessKeySize(sum - lower + 1);
            long num2 = map.lessKeySize(sum - up);
            ans += num1 - num2;
            map.put(sum);
        }
        return ans;
    }

    public static class SizeBalancedTreeMap{
        private static class SBTNode{
            private long key;
            private long size;
            private SBTNode l;
            private SBTNode r;
            private long num;

            public SBTNode(long key){
                this.key = key;
                size = 1;
                l = null;
                r = null;
                num = 1;
            }
        }

        private SBTNode root;

        private SBTNode rightRotate(SBTNode cur){
            long same = cur.num - (cur.l == null ? 0 : cur.l.num) - (cur.r == null ? 0 : cur.r.num);
            SBTNode left = cur.l;
            cur.l = left.r;
            left.r = cur;
            left.size = cur.size;
            cur.size = (cur.l == null ? 0 : cur.l.size) + (cur.r == null ? 0 : cur.r.size) + 1;
            left.num = cur.num;
            cur.num = (cur.r == null ? 0 : cur.r.num) + (cur.l == null ? 0 : cur.l.num) + same;
            return left;
        }

        private SBTNode leftRotate(SBTNode cur){
            long same = cur.num - (cur.l == null ? 0 : cur.l.num) - (cur.r == null ? 0 : cur.r.num);
            SBTNode right = cur.r;
            cur.r = right.l;
            right.l = cur;
            right.size = cur.size;
            cur.size = (cur.l == null ? 0 : cur.l.size) + (cur.r == null ? 0 : cur.r.size) + 1;
            right.num = cur.num;
            cur.num = (cur.r == null ? 0 : cur.r.num) + (cur.l == null ? 0 : cur.l.num) + same;
            return right;
        }

        private SBTNode maintain(SBTNode cur){
            if(cur == null)
                return null;
            long leftSize = cur.l == null ? 0 : cur.l.size;
            long rightSize = cur.r == null ? 0 : cur.r.size;
            long leftLeftSize = cur.l == null ? 0 : cur.l.l == null ? 0 : cur.l.l.size;
            long leftRightSize = cur.l == null ? 0 : cur.l.r == null ? 0 : cur.l.r.size;
            long rightRightSize = cur.r == null ? 0 : cur.r.r == null ? 0 : cur.r.r.size;
            long rightLeftSize = cur.r == null ? 0 : cur.r.l == null ? 0 : cur.r.l.size;
            if(rightSize < leftLeftSize){
                cur = rightRotate(cur);
                cur.r = maintain(cur.r);
                cur = maintain(cur);
            }else if(leftRightSize > rightSize){
                cur.l = leftRotate(cur.l);
                cur = rightRotate(cur);
                cur.l = maintain(cur.l);
                cur.r = maintain(cur.r);
                cur = maintain(cur);
            }else if(rightRightSize > leftSize){
                cur = leftRotate(cur);
                cur.l = maintain(cur.l);
                cur = maintain(cur);
            }else if(rightLeftSize > leftSize){
                cur.r = rightRotate(cur.r);
                cur = leftRotate(cur);
                cur.l = maintain(cur.l);
                cur.r = maintain(cur.r);
                cur =maintain(cur);
            }
            return cur;
        }

        private SBTNode add(SBTNode cur, long key){
            if(cur == null)
                return new SBTNode(key);
            if(cur.key > key){
                cur.l = add(cur.l, key);
            }else
                cur.r = add(cur.r, key);
            cur.size++;
            cur.num++;
            return maintain(cur);
        }

        private void hasContainProcess(long key){
            SBTNode cur = root;
            while(cur.key != key){
                cur.num++;
                if(cur.key > key)
                    cur = cur.l;
                else
                    cur = cur.r;
            }
            cur.num++;
        }

        private SBTNode findLastIndex(long key){
            SBTNode cur = root;
            SBTNode pre = root;
            while(cur != null){
                pre = cur;
                if(cur.key == key)
                    return cur;
                if(cur.key > key)
                    cur = cur.l;
                else
                    cur = cur.r;
            }
            return pre;
        }

        public void put(long key){
            SBTNode lastIndex = findLastIndex(key);
            if(lastIndex != null && lastIndex.key == key){
                hasContainProcess(key);
            }else {
                root = add(root, key);
            }
        }

        public long lessKeySize(long key){
            SBTNode cur = root;
            long sum = 0;
            while(cur != null){
                if(key == cur.key){
                   return sum + (cur.l == null ? 0 : cur.l.num);
                }
                if(key > cur.key){
                    sum += cur.num - (cur.r == null ? 0 : cur.r.num);
                    cur = cur.r;
                }else {
                    cur = cur.l;
                }
            }
            return sum;
        }

    }

    public static int countRangeSum1(int[] nums, int lower, int upper) {
        int n = nums.length;
        long[] sums = new long[n + 1];
        for (int i = 0; i < n; ++i)
            sums[i + 1] = sums[i] + nums[i];
        return countWhileMergeSort(sums, 0, n + 1, lower, upper);
    }

    private static int countWhileMergeSort(long[] sums, int start, int end, int lower, int upper) {
        if (end - start <= 1)
            return 0;
        int mid = (start + end) / 2;
        int count = countWhileMergeSort(sums, start, mid, lower, upper)
                + countWhileMergeSort(sums, mid, end, lower, upper);
        int j = mid, k = mid, t = mid;
        long[] cache = new long[end - start];
        for (int i = start, r = 0; i < mid; ++i, ++r) {
            while (k < end && sums[k] - sums[i] < lower)
                k++;
            while (j < end && sums[j] - sums[i] <= upper)
                j++;
            while (t < end && sums[t] < sums[i])
                cache[r++] = sums[t++];
            cache[r] = sums[i];
            count += j - k;
        }
        System.arraycopy(cache, 0, sums, start, t - start);
        return count;
    }
    // for test
    public static void printArray(int[] arr) {
        for (int i = 0; i < arr.length; i++) {
            System.out.print(arr[i] + " ");
        }
        System.out.println();
    }

    // for test
    public static int[] generateArray(int len, int varible) {
        int[] arr = new int[len];
        for (int i = 0; i < arr.length; i++) {
            arr[i] = (int) (Math.random() * varible);
        }
        return arr;
    }

    public static void main(String[] args) {
        int len = 200;
        int varible = 50;
        for (int i = 0; i < 10000; i++) {
            int[] test = generateArray(len, varible);
            int lower = (int) (Math.random() * varible) - (int) (Math.random() * varible);
            int upper = lower + (int) (Math.random() * varible);
            int ans1 = countRangeSum1(test, lower, upper);
            int ans2 = countRangeSum2(test, lower, upper);
            if (ans1 != ans2) {
                printArray(test);
                System.out.println(lower);
                System.out.println(upper);
                System.out.println(ans1);
                System.out.println(ans2);
            }
        }
    }
}
